import os
import random
L_train = []
L_val = []
L_test = []
# 读取文件中的内容，并写入列表FileNameList
def ReadFileDatas(original_filename):
      FileNameList = []
      file = open(original_filename, 'r+', encoding='utf-8-sig')
      for line in file:
            FileNameList.append(line)  # 写入文件内容到列表中去
      print('数据集总量：', len(FileNameList))
      file.close()
      return FileNameList
# 将获取的列表中的内容转为 str ，再写入到txt文件中去
# listInfo为 ReadFileDatas 的列表
def WriteDatasToFile(listInfo, new_filename):
      file_handle = open(new_filename, mode='a', encoding='utf-8-sig')
      for idx in range(len(listInfo)):
            str = listInfo[idx]  # 列表指针
            str_Result = str
            file_handle.write(str_Result)
      file_handle.close()
      print('写入 %s 文件成功.' % new_filename)
      def TrainValTestFile(new_filename):
            i = 0  # counter
            j = 9352  # all lines
            file_divide = open(new_filename, 'r', encoding='utf-8-sig')
            lines = file_divide.readlines()
            for line in lines:
                  if i < (j * 0.6):
                        i += 1
                        L_train.append(line)
                  elif i < (j * 0.8):
                        i += 1
                        L_val.append(line)
                  elif i < j:
                        i += 1
                        L_test.append(line)
            print("总数据量：%d , 此时创建train, val, test数据集" % i)
            return L_train, L_val, L_test
      def text_save(filename, data):  # filename为写入CSV文件的路径，data为要写入数据列表
            file = open(filename, 'a', encoding='utf-8-sig')
            for i in range(len(data)):
                  s = str(data[i]).replace('[', '').replace(']', '')
                  file.write(s)
            file.close()
            print("保存数据集（路径）成功：%s" % filename)
            if __name__ == "__main__":
                  listFileInfo = ReadFileDatas('data.txt')
                  random.shuffle(listFileInfo)
                  WriteDatasToFile(listFileInfo, 'new_data.txt')
                  TrainValTestFile('new_data.txt')
                  text_save('./data/data_train.txt', L_train)
                  text_save('./data/data_val.txt', L_val)
                  text_save('./data/data_test.txt', L_test)